package com.ruoyi.web.controller.ai;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.swagger.annotations.Api;
import io.swagger.annotations.ApiOperation;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.RequestBody;
import org.apache.commons.lang3.StringEscapeUtils;
import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.MediaType;
import org.springframework.http.client.ClientHttpResponse;
import org.springframework.http.client.HttpComponentsClientHttpRequestFactory;
import org.springframework.http.client.SimpleClientHttpRequestFactory;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.client.ResponseExtractor;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

@RestController
@RequestMapping("ai")
@Api("大模型")
public class AiController {
    private static final Logger log = LoggerFactory.getLogger(AiController.class);

    ExecutorService executorService = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors());

    @ApiOperation("大模型工作流")
    @GetMapping("workflow")
    public ResponseBodyEmitter workflow() {
        // 设置超时时间为60秒
        ResponseBodyEmitter emitter = new ResponseBodyEmitter(60 * 1000L);

        // 开启异步线程处理数据并发送
        executorService.execute(() -> {
            try {
                for (int i = 0; i < 5; i++) {
                    String msg = String.format("%d,%d \n", System.currentTimeMillis(), i + 1);
                    log.info("{}", msg);
                    emitter.send(msg);
                    Thread.sleep(1000);
                }
                emitter.complete();
            } catch (Exception e) {
                emitter.completeWithError(e);
            }
        });
        return emitter;
    }

    @ApiOperation("sse")
    @GetMapping("sse")
    public SseEmitter handleSSE() {
        SseEmitter emitter = new SseEmitter();

        executorService.execute(() -> {
            try {
                // 模拟发送10条SSE事件
                for (int i = 0; i < 10; i++) {
                    emitter.send(SseEmitter.event()
                            .id(String.valueOf(i))
                            .name("message")
                            .data("SSE数据项 " + i));
                    Thread.sleep(1000);
                }
                emitter.complete();
            } catch (Exception e) {
                emitter.completeWithError(e);
            }
        });

        return emitter;
    }

    private final ObjectMapper objectMapper = new ObjectMapper();

    @ApiOperation("dify")
    @PostMapping("dify")
    public SseEmitter difyRequest(@RequestParam("query") String query) throws JsonProcessingException {

        SseEmitter emitter = new SseEmitter(120 * 1000L);

        Map<String, Object> params = Map.of(
                "user", "zxy",
                "response_mode", "streaming",
                "inputs", Map.of("query", query)
        );
        String s = objectMapper.writeValueAsString(params);
        log.info("{}", s);

        OkHttpClient client = new OkHttpClient.Builder()
                .callTimeout(5, TimeUnit.MINUTES)
                .connectTimeout(5, TimeUnit.MINUTES)
                .readTimeout(5, TimeUnit.MINUTES)
                .writeTimeout(5, TimeUnit.MINUTES)
                .build();
        Request request = new Request.Builder()
                .url("http://localhost/v1/workflows/run")
                .post(RequestBody.create(s.getBytes()))
                .addHeader("Authorization", "Bearer " + "app-DXjZ5Ds57ARiZl7PC4HFFnr6")
                .addHeader("Content-Type", "application/json")
                .build();
        client.newCall(request).enqueue(new okhttp3.Callback() {
            @Override
            public void onResponse(@NotNull okhttp3.Call call, @NotNull okhttp3.Response response) {
                try (okhttp3.ResponseBody body = response.body();
                     BufferedReader reader = new BufferedReader(new InputStreamReader(body.byteStream()))) {
                    String line;
                    while ((line = reader.readLine()) != null) {
                        if (!line.trim().isEmpty()) {
                            String value = StringEscapeUtils.unescapeJava(line);
                            emitter.send(value);
                        }
                    }
                    emitter.complete();

                } catch (Exception e) {
                    emitter.completeWithError(e);
                }
            }

            @Override
            public void onFailure(@NotNull okhttp3.Call call, @NotNull IOException e) {
                emitter.completeWithError(e);
            }

        });
        return emitter;
    }

}